from sqlalchemy import Column, String, Integer
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from init.mysql import engine

Base = declarative_base()


class StockIncrease(Base):
    __tablename__ = 'stock_increase'
    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(String(6))
    sdate = Column(String(10))
    type = Column(String(50))
    price = Column(String(20))
    tprice = Column(String(20))
    fprice = Column(String(20))
    amount = Column(String(20))

    def __init__(self, code, sdate, type, price, tprice, fprice, amount):
        self.code = code
        self.sdate = sdate
        self.type = type
        self.price = price
        self.tprice = tprice
        self.fprice = fprice
        self.amount = amount

    @staticmethod
    def insert_or_update(code, data_list):
        session = sessionmaker(bind=engine)()
        for data in data_list:
            existing_stock_increase = session.query(StockIncrease).filter_by(code=code, sdate=data['sdate']).first()
            if existing_stock_increase:
                existing_stock_increase.sdate = data['sdate']
                existing_stock_increase.type = data['type']
                existing_stock_increase.price = data['price']
                existing_stock_increase.tprice = data['tprice']
                existing_stock_increase.fprice = data['fprice']
                existing_stock_increase.amount = data['amount']

            else:
                session.add(StockIncrease(code=code, sdate=data['sdate'], type=data['type'], price=data['price'],
                                          tprice=data['tprice'], fprice=data['fprice'], amount=data['amount']))
        session.commit()


Base.metadata.create_all(engine)
